# Author: LIAN DONGJIE
# Date: 2024-12-28
# Description: A mono-window algorithm for retrieving land surface temperature from Landsat TM5 data

import os
import numpy as np
import matplotlib.pyplot as plt
from osgeo import gdal
import math
from matplotlib.colors import LinearSegmentedColormap

class LSTRetrieval:
    """
    单窗算法实现Landsat TM5地表温度反演
    """
    def __init__(self, folder_path):
        """
        初始化类实例
        
        参数:
            folder_path: 包含Landsat TM5数据的文件夹路径
        """
        self.folder = os.path.basename(folder_path)
        self.prefix = self.folder
        self.folder_path = folder_path
        
        # 初始化数据存储变量
        self.datasets = {}
        self.dn_values = {}
        self.radiance = {}
        self.ndvi = None
        self.fvc = None
        self.emissivity = None
        self.surface_temperature = None
        
        # 可视化参数
        self.figsize = (12, 10)
        self.cmap_ndvi = LinearSegmentedColormap.from_list('ndvi', ['red', 'white', 'green'])
        self.cmap_temp = plt.cm.jet
        
    def read_data(self):
        """读取Landsat TM5各波段数据"""
        bands = {
            'red': 'B3',
            'green': 'B2',
            'blue': 'B1',
            'nir': 'B4',
            'thermal': 'B6'
        }
        
        for band_name, band_id in bands.items():
            file_path = os.path.join(self.folder_path, f"{self.prefix}_{band_id}.tif")
            self.datasets[band_name] = gdal.Open(file_path)
            
            if self.datasets[band_name] is None:
                raise FileNotFoundError(f"无法打开文件: {file_path}")
                
            # 读取DN值并转换为float64
            self.dn_values[band_name] = np.array(
                self.datasets[band_name].GetRasterBand(1).ReadAsArray()
            ).astype(np.float64)
            
            # 将DN=0的像素设为NaN
            self.dn_values[band_name][self.dn_values[band_name] == 0] = np.nan
            
        print("数据读取完成")
        
    def read_calibration_coefficients(self):
        """读取辐射定标系数"""
        mtl_path = os.path.join(self.folder_path, f"{self.prefix}_MTL.txt")
        
        if not os.path.exists(mtl_path):
            raise FileNotFoundError(f"元数据文件不存在: {mtl_path}")
            
        with open(mtl_path, 'r') as file:
            lines = file.readlines()
            
        # 找到辐射定标参数的起始行
        for i in range(len(lines)):
            if 'GROUP = RADIOMETRIC_RESCALING' in lines[i]:
                break
                
        self.calibration_start = i
        
        def get_gain_offset(band_index):
            """获取特定波段的增益和偏移量"""
            gain = float(lines[i + band_index][:-1].split(' ')[-1])
            offset = float(lines[i + band_index + 7][:-1].split(' ')[-1])
            return (gain, offset)
            
        # 获取各波段的增益和偏移量
        self.gains = {}
        self.offsets = {}
        
        band_indices = {
            'red': 3,
            'green': 2,
            'blue': 1,
            'nir': 4,
            'thermal': 6
        }
        
        for band_name, index in band_indices.items():
            self.gains[band_name], self.offsets[band_name] = get_gain_offset(index)
            
        print("辐射定标系数读取完成")
        
    def radiometric_calibration(self):
        """执行辐射定标，将DN值转换为辐射亮度"""
        for band_name in self.dn_values:
            if band_name == 'thermal':
                # 热红外波段不需要乘以0.1
                self.radiance[band_name] = (
                    self.gains[band_name] * self.dn_values[band_name] + 
                    self.offsets[band_name]
                )
            else:
                # 其他波段需要乘以0.1
                self.radiance[band_name] = 0.1 * (
                    self.gains[band_name] * self.dn_values[band_name] + 
                    self.offsets[band_name]
                )
                
            # 将辐射亮度小于等于0的像素设为NaN
            self.radiance[band_name][self.radiance[band_name] <= 0] = np.nan
            
        print("辐射定标完成")
        
    def calculate_ndvi(self):
        """计算归一化植被指数(NDVI)"""
        if 'nir' not in self.radiance or 'red' not in self.radiance:
            raise ValueError("需要先完成辐射定标才能计算NDVI")
            
        self.ndvi = (self.radiance['nir'] - self.radiance['red']) / (
            self.radiance['nir'] + self.radiance['red']
        )
        
        print(f"NDVI计算完成，范围: {np.nanmin(self.ndvi)} - {np.nanmax(self.ndvi)}")
        
    def calculate_fvc(self):
        """计算植被覆盖度(FVC)"""
        if self.ndvi is None:
            raise ValueError("需要先计算NDVI才能计算FVC")
            
        # 计算5%和95%分位数作为土壤和植被的NDVI阈值
        ndvi_min, ndvi_max = np.nanpercentile(self.ndvi, 5), np.nanpercentile(self.ndvi, 95)
        
        # 计算植被覆盖度
        self.fvc = np.where(
            self.ndvi < ndvi_min, 0,
            np.where(
                self.ndvi > ndvi_max, 1,
                (self.ndvi - ndvi_min) / (ndvi_max - ndvi_min)
            )
        )
        
        print(f"植被覆盖度计算完成，范围: {np.nanmin(self.fvc)} - {np.nanmax(self.fvc)}")
        
    def calculate_emissivity(self):
        """计算地表比辐射率"""
        if self.fvc is None:
            raise ValueError("需要先计算FVC才能计算地表比辐射率")
            
        # 使用Sobrino et al. (2004)的方法计算地表比辐射率
        self.emissivity = (
            self.fvc * 0.985 + 
            (1 - self.fvc) * 0.960 + 
            0.06 * self.fvc * (1 - self.fvc)
        )
        
        print(f"地表比辐射率计算完成，范围: {np.nanmin(self.emissivity)} - {np.nanmax(self.emissivity)}")
        
    def calculate_surface_temperature(self, air_temperature=15, atmospheric_transmittance=0.7):
        """
        计算地表温度
        
        参数:
            air_temperature: 大气温度(°C)
            atmospheric_transmittance: 大气透过率
        """
        if self.emissivity is None:
            raise ValueError("需要先计算地表比辐射率才能计算地表温度")
            
        # 将大气温度转换为开尔文
        T0 = air_temperature + 273.15
        
        # 计算平均大气作用温度
        Ta = 16.0110 + 0.92621 * T0
        
        # 计算热红外波段的亮温
        K1, K2 = 607.76, 1260.56  # Landsat TM5热红外波段的校准常数
        T_thermal = K2 / np.log(1 + K1 / self.radiance['thermal'])
        
        # 单窗算法参数
        C = self.emissivity * atmospheric_transmittance
        D = (1 - atmospheric_transmittance) * (1 + (1 - self.emissivity) * atmospheric_transmittance)
        
        # 单窗算法系数(a和b)
        a, b = -67.355351, 0.458606
        
        # 计算地表温度(开尔文)
        Ts_kelvin = ((a * (1 - C - D) + (b * (1 - C - D) + C + D) * T_thermal - D * Ta) / C)
        
        # 转换为摄氏度
        self.surface_temperature = Ts_kelvin - 273.15
        
        print(f"地表温度计算完成，范围: {np.nanmin(self.surface_temperature)} - {np.nanmax(self.surface_temperature)}°C")
        
    def visualize_results(self, output_dir=None):
        """
        可视化中间结果和最终地表温度
        
        参数:
            output_dir: 图像保存目录，如果为None则不保存
        """
        if self.surface_temperature is None:
            raise ValueError("需要先计算地表温度才能进行可视化")
            
        plt.figure(figsize=self.figsize)
        
        # 1. 显示NDVI
        plt.subplot(2, 2, 1)
        plt.imshow(self.ndvi, cmap=self.cmap_ndvi, vmin=-1, vmax=1)
        plt.colorbar(label='NDVI')
        plt.title('归一化植被指数 (NDVI)')
        plt.axis('off')
        
        # 2. 显示植被覆盖度
        plt.subplot(2, 2, 2)
        plt.imshow(self.fvc, cmap='viridis', vmin=0, vmax=1)
        plt.colorbar(label='FVC')
        plt.title('植被覆盖度 (FVC)')
        plt.axis('off')
        
        # 3. 显示地表比辐射率
        plt.subplot(2, 2, 3)
        plt.imshow(self.emissivity, cmap='plasma', vmin=0.9, vmax=1.0)
        plt.colorbar(label='Emissivity')
        plt.title('地表比辐射率')
        plt.axis('off')
        
        # 4. 显示地表温度
        plt.subplot(2, 2, 4)
        temp_plot = plt.imshow(self.surface_temperature, cmap=self.cmap_temp)
        plt.colorbar(label='Temperature (°C)')
        plt.title('地表温度 (°C)')
        plt.axis('off')
        
        plt.tight_layout()
        
        if output_dir:
            os.makedirs(output_dir, exist_ok=True)
            output_path = os.path.join(output_dir, f"{self.prefix}_LST_results.png")
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            print(f"可视化结果已保存至: {output_path}")
            
        plt.show()
        
    def save_surface_temperature(self, output_path=None):
        """
        保存地表温度结果为GeoTIFF文件
        
        参数:
            output_path: 输出文件路径，如果为None则使用默认路径
        """
        if self.surface_temperature is None:
            raise ValueError("需要先计算地表温度才能保存结果")
            
        if output_path is None:
            output_path = os.path.join(self.folder_path, f"{self.prefix}_LST.tif")
            
        # 获取原始热红外波段的信息
        thermal_dataset = self.datasets['thermal']
        width = thermal_dataset.RasterXSize
        height = thermal_dataset.RasterYSize
        bands = thermal_dataset.RasterCount
        datatype = gdal.GDT_Float32
        
        # 获取地理变换和投影信息
        geotransform = thermal_dataset.GetGeoTransform()
        projection = thermal_dataset.GetProjection()
        
        # 创建新的GeoTIFF文件
        driver = gdal.GetDriverByName('GTiff')
        dst_dataset = driver.Create(output_path, width, height, bands, datatype)
        
        # 设置地理变换和投影信息
        dst_dataset.SetGeoTransform(geotransform)
        dst_dataset.SetProjection(projection)
        
        # 写入地表温度数据
        dst_band = dst_dataset.GetRasterBand(1)
        dst_band.WriteArray(self.surface_temperature)
        
        # 关闭数据集
        del dst_dataset
        
        print(f"地表温度结果已保存至: {output_path}")
        
    def run_full_process(self, air_temperature=15, atmospheric_transmittance=0.7, 
                         visualize=True, output_dir=None, save_result=True):
        """
        运行完整的地表温度反演流程
        
        参数:
            air_temperature: 大气温度(°C)
            atmospheric_transmittance: 大气透过率
            visualize: 是否可视化结果
            output_dir: 图像和结果保存目录
            save_result: 是否保存地表温度结果
        """
        self.read_data()
        self.read_calibration_coefficients()
        self.radiometric_calibration()
        self.calculate_ndvi()
        self.calculate_fvc()
        self.calculate_emissivity()
        self.calculate_surface_temperature(air_temperature, atmospheric_transmittance)
        
        if visualize:
            self.visualize_results(output_dir)
            
        if save_result:
            output_path = os.path.join(output_dir, f"{self.prefix}_LST.tif") if output_dir else None
            self.save_surface_temperature(output_path)


if __name__ == "__main__":
    # 我一般是一个工程，里面包含input,output,process3个文件夹
    folder_path = 'C:/Users/admin/Desktop/Project/input'
    lst_processor = LSTRetrieval(folder_path)
    
    
    lst_processor.run_full_process(
        air_temperature=15,  # 大气温度(°C)
        atmospheric_transmittance=0.7,  # 大气透过率
        visualize=True,  # 可视化结果
        output_dir='./output',  # 结果保存目录
        save_result=True  # 保存地表温度结果
    )
    
